# coding:utf-8
# Author : hiicy redldw
# Date : 2019/02/25
import tensorflow as tf
from tensorflow.contrib import rnn
import numpy as np
import input_data
import warnings

warnings.filterwarnings('ignore', '.*')


def MNISTRECOG():
    mnist = input_data.read_data_sets('MNIST_data', one_hot=True)
    # O * W + b -> 10 labels for each image, O[? 28], W[28 10], B[10]
    # FIXME:time_step_size表示的是lstm神经元的个数，这里设置为28个(和图片的尺寸有关？)，input_vec_size表示一次输入的像素数
    input_vec_size = lstm_size = 28  # 输入向量维度
    time_step_size = 28  # 循环层输出
    batch_size = 128
    test_size = 256

    def init_weights(shape):
        return tf.Variable(tf.random_normal(shape, stddev=0.01))

    def model(X, W, B, lstm_size):
        # XT shape: (time_step_size, batch_size, input_vec_size)
        # FAQ:对这一步操作还不是太理解，为什么需要将第一行和第二行置换
        XT = tf.transpose(X, [1, 0, 2])
        # XR shape: (time_step_size * batch_size, input_vec_size)
        XR = tf.reshape(XT, [-1, lstm_size])
        # Each array shape: (batch_size, input_vec_size)
        X_split = tf.split(XR, time_step_size,
                           0)  # split them to time_step_size (28 arrays),shape = [(128, 28),(128, 28)...]
        #  Make lstm with lstm_size (each input vector size) num_units=lstm_size
        lstm = rnn.BasicLSTMCell(lstm_size, forget_bias=1.0, state_is_tuple=True)

        # Get lstm cell output, time_step_size (28) arrays with lstm_size output: (batch_size, lstm_size)
        # rnn..static_rnn()的输出对应于每一个timestep，如果只关心最后一步的输出，取outputs[-1]即可
        outputs, _states = rnn.static_rnn(lstm, X_split, dtype=tf.float32)  # 时间序列上每个Cell的输出:[... shape=(128, 28)..]
        # Get the last output
        return tf.matmul(outputs[-1], W) + B, lstm.state_size  # State size to initialize the state

    trX, trY, teX, teY = mnist.train.images, mnist.train.labels, mnist.test.images, mnist.test.images
    # mnist.train.images是一个55000 * 784维的矩阵, mnist.train.labels是一个55000 * 10维的矩阵
    # 训练集包含55000张图片，每张图片为28*28维矩阵
    trX = trX.reshape((-1, 28, 28))
    teX = teX.reshape((-1, 28, 28))
    print(trX.shape, trX.dtype, trY.dtype)
    X = tf.placeholder(tf.float32, shape=[None, 28, 28])
    Y = tf.placeholder(tf.float64, shape=[None, 10])

    W = init_weights([lstm_size, 10])  # # 输出层权重矩阵28×10
    B = init_weights([10])  ## 输出层bais

    py_x, state_size = model(X, W, B, lstm_size)
    cost = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=py_x, labels=Y))
    train_op = tf.train.RMSPropOptimizer(0.001, 0.9).minimize(cost)

    # 返回每一行的最大值
    predict_op = tf.argmax(py_x, 1)
    # tf.ConfigProto，一般是在创建session时对session进行配置
    session_conf = tf.ConfigProto()
    session_conf.gpu_options.allow_growth = True  # 允许gpu在使用的过程中慢慢增加。
    with tf.Session(config=session_conf) as sess:
        tf.global_variables_initializer().run()
        for i in range(50):
            # 从训练集中每段选择一个batch训练，batch_size= end-start
            for start, end in zip(range(0, len(trX), batch_size), range(batch_size, len(trX) + 1, batch_size)):
                sess.run(train_op, feed_dict={X: trX[start:end], Y: trY[start:end]})
                # X (128,28,28)
            s = len(teX)
            test_indices = np.arange(len(teX))  # Get A Test Batch
            np.random.shuffle(test_indices)
            test_indices = test_indices[0:test_size]
            print(i, np.mean(np.argmax(teY[test_indices], axis=1) ==
                             sess.run(predict_op, feed_dict={X: teX[test_indices]})))





